#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include "CFG.h"
#include "../DynArray.h"
#include "TU.h"
#include "BasicBlock.h"

//CFGEdge

CFGEdge* CFGE_New(CFGNode* TargetNode){
    CFGEdge* New;
    
    New = malloc(sizeof(CFGEdge));
    New->Count = 0;
    New->Target = TargetNode;
    
    return New;
}

//CFGNode

static int EdgeTargetComparator(void* E1, void* E2){
    CFGEdge *Edge1, *Edge2;
    
    Edge1 = E1;
    Edge2 = E2;
    
    return Edge1->Target == Edge2->Target ?  1 : 0;
}

static int SourceNodeComparator(void* N1, void *N2){
    CFGNode *Node1, *Node2;
    
    return Node1 == Node2 ? 1 : 0;
}

CFGNode* CFGN_New(TranslationUnit* TU){
    CFGNode* New;
    
    New = malloc(sizeof(CFGNode));
    
    New->Unit = TU;
    New->Edges = LST_New(EdgeTargetComparator);
    New->_Visited = 0;
    New->EpochExecutionCount = 0;
    New->TARJ_Index = 0;
    New->TARJ_LowIndex = 0;
    New->SCC_Index = 0;
    New->TARJ_HighestTransitionCount = 0;
    New->SCC_HighestTransitionCount = 0;
    New->SCC_HotNodes = NULL;
    New->SourceNodes = LST_New(SourceNodeComparator);
    New->CYC_Color = DFSC_WHITE;
    New->Size = 0;
    
    return New;
}

void CFGN_AddEdge(CFGNode* CFGN, CFGEdge* Edge){
    LST_Enqueue(CFGN->Edges, Edge);
}

void* CFGN_RemoveEdge(CFGNode* CFGN, CFGNode* Target){
    return LST_Remove(CFGN->Edges, Target);
}

CFGEdge* CFGN_EdgeTo(CFGNode* CFGN, uint16_t Address){
    uint32_t Count, i;
    CFGEdge *Edge, *Ret = NULL;
    
    LST_Rewind(CFGN->Edges);
    
    while((Edge = LST_Next(CFGN->Edges))){
        if(TU_GetStartingAddress(Edge->Target->Unit) == Address){
            Ret = Edge;
            break;
        }
    }
    
    #ifdef TRANSITION_DEBUG
    printf("Ret = %p\n", Ret);
    #endif
    
    return Ret;
}

static void CFGN_AddSourceNode(CFGNode* DestNode, CFGNode* SourceNode){
    CFGNode* Node;
    LST_Rewind(DestNode->SourceNodes);
    while((Node = LST_Next(DestNode->SourceNodes))){
        if(Node == SourceNode)
            return;
    }
    LST_Insert(DestNode->SourceNodes, SourceNode);
}

static void CFGN_FixSourceNode(CFGNode* DestNode, CFGNode* OldSource, CFGNode* NewSource){
    LST_Remove(DestNode->SourceNodes, OldSource);
    LST_Insert(DestNode->SourceNodes, NewSource);
}

//CFG Node Map

CFGNodeMap*  CFGNMNew(){
    CFGNodeMap* New;
    
    New = malloc(sizeof(CFGNodeMap));  
    New->Data = NULL;
    New->Children = malloc(sizeof(CFGNodeMap*) * 16);
    memset(New->Children, 0, 16 * sizeof(CFGNodeMap*));    
    New->MapKey = 0;
    
    return New;
}

static uint32_t Hash(uint16_t Chave){
    return (Chave & 0xF);
}

static void CFGNMInsert_(CFGNodeMap* NM, CFGNode* Node, uint32_t Depth){
    uint32_t Pos;
    uint16_t Key;
    
    if(!(NM->Data)){
        NM->Data = Node;
        NM->MapKey = TU_GetStartingAddress(Node->Unit);
    }
    else{
        Key = TU_GetStartingAddress(Node->Unit);
        Pos = Hash(Key >> (Depth * 4));
        
        if(!NM->Children[Pos])
            NM->Children[Pos] = CFGNMNew();
        
        CFGNMInsert_(NM->Children[Pos], Node, Depth + 1);
    }
}

static void CFGNMInsert(CFGNodeMap* NM, CFGNode* Node){
    uint32_t Pos;
    uint16_t Key;
    
    Key = TU_GetStartingAddress(Node->Unit);
    Pos = Hash(Key);
    
    if(!NM->Children[Pos])
        NM->Children[Pos] = CFGNMNew();
    
    CFGNMInsert_(NM->Children[Pos], Node, 1);
}

static CFGNode* CFGNMSearch_(CFGNodeMap *NM, uint16_t StartingAddress, uint32_t Depth){
    uint32_t Pos;
    
    if(NM->Data){
        if(NM->MapKey == StartingAddress){
            return NM->Data;
        }
        else{
            Pos = Hash(StartingAddress >> (Depth * 4));
            if(NM->Children[Pos])
                CFGNMSearch_(NM->Children[Pos], StartingAddress, Depth + 1);
            else
                return NULL;
        }
    }
    else
        return NULL;
}

static CFGNode* CFGNMSearch(CFGNodeMap *NM, uint16_t StartingAddress){
    uint32_t Pos;
    
    Pos = Hash(StartingAddress);
    
    return NM->Children[Pos] ? 
        CFGNMSearch_(NM->Children[Pos], StartingAddress, 1) 
            : NULL;    
}

static CFGNode* CFGNMUpdate_(CFGNodeMap *NM, uint16_t StartingAddress, CFGNode* NewData, uint32_t Depth){
    uint32_t Pos;
 
    if(NM->Data){
        if(NM->MapKey == StartingAddress){
            NM->Data = NewData;
            return NM->Data;
        }
        else{
            Pos = Hash(StartingAddress >> (Depth * 4));
            if(NM->Children[Pos])
                CFGNMUpdate_(NM->Children[Pos], StartingAddress, NewData ,Depth + 1);
            else
                return NULL;
        }
    }
    else
        return NULL;    
}

//This function DOES NOT update MapKey! This makes sense since we don't want to
//change the trie's indices.
static CFGNode* CFGNMUpdate(CFGNodeMap *NM, uint16_t StartingAddress, CFGNode* NewData){
    uint32_t Pos;
    
    Pos = Hash(StartingAddress);
    
    return CFGNMUpdate_(NM->Children[Pos], StartingAddress, NewData, 1);       
}

//CFG

static int CompareCFGNodes(void* Node1, void* Node2){
    CFGNode *N1, *N2;
    N1 = Node1;
    N2 = Node2;
    
    return TU_GetStartingAddress(N1->Unit) - TU_GetStartingAddress(N2->Unit);
}

ControlFlowGraph* CFG_New(void){
    ControlFlowGraph* New;
    
    New = malloc(sizeof(ControlFlowGraph));
    
    New->InitialAddress = 0;
    New->Nodes = CFGNMNew();
    New->NodeList = LST_New(CompareCFGNodes);
    New->NumNodes = 0;
    
    return New;
}

//Simply adds a node to the CFG node list and updates counter
void CFG_AddNewNode(ControlFlowGraph* CFG, CFGNode* Node){
    //Set the initial address to the starting address of the first node included
    if(!CFG->NumNodes)
        CFG->InitialAddress = TU_GetStartingAddress(Node->Unit);
    
    LST_Insert(CFG->NodeList, Node);
    CFG->NumNodes++;
}

//Adds a node to the CFG and updates de CFG Node Map and CFG Repository
void CFG_AddNode(ControlFlowGraph* CFG, CFGNode* Node){
    
    CFG_AddNewNode(CFG, Node);
    
    //Insert node into the CFG node map
    CFGNMInsert(CFG->Nodes, Node);
    
    //Refresh CFG Repository (for current node)
    CFGRNodeRefresh(CFG, Node);
}

static void CFGUpdateCFGNodes(ControlFlowGraph* CFG, ControlFlowGraph *CFGN, CFGNode* NewCFGNode){
    CFGNode* Node;
    
    LST_Rewind(CFGN->NodeList);
    while((Node = LST_Next(CFGN->NodeList))){
        if(TU_GetBB(Node->Unit))
            CFGNMUpdate(CFG->Nodes, TU_GetStartingAddress(Node->Unit), NewCFGNode);
        else
            CFGUpdateCFGNodes(CFG, TU_GetCFG(Node->Unit), NewCFGNode);
    }
}

static void CFG_UpdateNodes(ControlFlowGraph* CFG, CFGNode* NewCFGNode){
    CFGNode *Node;
    ControlFlowGraph *NewCFG;
    
    CFG_AddNewNode(CFG, NewCFGNode);
    
    NewCFG = TU_GetCFG(NewCFGNode->Unit);
    
    LST_Rewind(NewCFG->NodeList);
    
    //Fix the node map to point to the new CFG instead of the old translation units
    //Also, remove the nodes from the parent CFG's node list.
    while((Node = LST_Next(NewCFG->NodeList))){
        
        //THIS IS A NASTY S.E. Hack. Fix this later
        //BEGIN HACK
        if(TU_GetBB(Node->Unit))
            CFGNMUpdate(CFG->Nodes, TU_GetStartingAddress(Node->Unit), NewCFGNode);
        else
            CFGUpdateCFGNodes(CFG, TU_GetCFG(Node->Unit), NewCFGNode);
        //END HACK
        
        if(TU_GetStartingAddress(Node->Unit) == CFG->InitialAddress){
            CFG->InitialAddress = TU_GetStartingAddress(NewCFGNode->Unit);
        }
        
        LST_Remove(CFG->NodeList, Node);
        CFG->NumNodes--;
    }
    
    //This may not be needed...
    //CFGRNodeRefresh(CFG, NewCFGNode);
}

CFGNode* CFG_SearchNode(ControlFlowGraph* CFG, uint16_t InitialAddress){
    return CFGNMSearch(CFG->Nodes, InitialAddress);
}

void CFG_AddBBNode(ControlFlowGraph* CFG, TranslationUnit* TU){
    CFGNode* NewCFGNode;
    NewCFGNode = CFGN_New(TU);
    CFG_AddNode(CFG, NewCFGNode);    
}


void CFG_AddCFGNode(ControlFlowGraph* CFG, TranslationUnit* TU){
    CFGNode* NewCFGNode;
    NewCFGNode = CFGN_New(TU);
    CFG_UpdateNodes(CFG, NewCFGNode);
}

void CFG_IncrementEdge(ControlFlowGraph* CFG, TranslationUnit* SourceTU, TranslationUnit* TargetTU){
    CFGNode* SourceNode, *TargetNode;
    CFGEdge* Edge;
    
    SourceNode = CFG_SearchNode(CFG, TU_GetStartingAddress(SourceTU));
    TargetNode = CFG_SearchNode(CFG, TU_GetStartingAddress(TargetTU));
    
    CFGN_AddSourceNode(TargetNode, SourceNode);
    
    Edge = CFGN_EdgeTo(SourceNode, TU_GetStartingAddress(TargetTU));
    
    if(Edge){
        Edge->Count++;
    }
    else{
        Edge = CFGE_New(TargetNode);
        Edge->Count = 1;
        CFGN_AddEdge(SourceNode, Edge);
    }
}

//This function is responsible for fusing a group of Nodes belonging to a CFG
//into a separate CFG, and then connecting it back to the original CFG
void CFG_Fuse(ControlFlowGraph* CFG, __LinkedList* Nodes){
    ControlFlowGraph *NewCFG;
    CFGNode *Node, *SourceNode, *NewCFGNode;
    TranslationUnit* CFGTU;
    CFGEdge* Edge, *NewEdge, *OldEdge;
    uint64_t NodeExecutions = 0;
    
    NewCFG = CFG_New();
    
    LST_Rewind(Nodes);
    
    //Add all nodes into the new CFG
    while((Node = LST_Next(Nodes))){
        CFG_AddNode(NewCFG, Node);
        NodeExecutions += Node->EpochExecutionCount;
        Node->EpochExecutionCount = 0;
    }
    
    printf("Fusing nodes into new CFG: 0x%04X\n", NewCFG->InitialAddress);
    
    //Add the new CFG as a node to the original CFG
    CFGTU = TU_New();
    TU_SetCFG(CFGTU, NewCFG);
    
    CFG_AddCFGNode(CFG, CFGTU);
    
    NewCFGNode = CFG_SearchNode(CFG, NewCFG->InitialAddress);
    
    NewCFGNode->EpochExecutionCount = NodeExecutions;
    
    //Fix all incoming and outgoing edges
    LST_Rewind(NewCFG->NodeList);
    while((Node = LST_Next(NewCFG->NodeList))){
        //Go through all outgoing edges
        LST_Rewind(Node->Edges);
        while((Edge = LST_Next(Node->Edges))){
            //Check if outgoing edge is pointing to a node outside the new CFG
            if(!CFG_SearchNode(NewCFG, TU_GetStartingAddress(Edge->Target->Unit))){
                //If so, fix it by adding an edge from the NewCFG Node into the target node
                
                NewEdge = CFGN_EdgeTo(NewCFGNode, TU_GetStartingAddress(Edge->Target->Unit));
                
                if(!NewEdge){
                    NewEdge = CFGE_New(Edge->Target);
                    CFGN_AddEdge(NewCFGNode, NewEdge);
                }
                
                NewEdge->Count = Edge->Count;
                
                //Fix SourceNode for the target
                CFGN_FixSourceNode(Edge->Target, Node, NewCFGNode);
                
                //Remove the original edge.
                free(CFGN_RemoveEdge(Node, Edge->Target));
                
            }
        }

        //Go through all incoming edges
        LST_Rewind(Node->SourceNodes);
        while((SourceNode = LST_Next(Node->SourceNodes))){
            //Check if incoming edge is from outside the new CFG
            if(!CFG_SearchNode(NewCFG, TU_GetStartingAddress(SourceNode->Unit))){
                //If so, fix it so it points to the new CFG node
                LST_Rewind(SourceNode->Edges);
                while((Edge = LST_Next(SourceNode->Edges))){
                    //if(TU_GetStartingAddress(Edge->Target->Unit) == TU_GetStartingAddress(Node->Unit)){
                    if(TU_IsSameTU(Edge->Target->Unit, Node->Unit)){
                        Edge->Target = NewCFGNode;
                    }
                }
            }
        }
    }
    
}

void CFG_UpdateNodeSizes(ControlFlowGraph* CFG){
    CFGNode *Node;
    
    LST_Rewind(CFG->NodeList);
    
    while((Node = LST_Next(CFG->NodeList))){
        Node->_VisitedSize = 0;
    }
    
    LST_Rewind(CFG->NodeList);
    
    while((Node = LST_Next(CFG->NodeList))){
        if(!Node->_VisitedSize){
            Node->Size = TU_GetSize(Node->Unit);
        }
    }
}

//This must be called every time a full CFG traversal is done.
//It resets all traversal variables. Couldn't think of a better
//way to do this.
void CFG_ResetCFGForTraversal(ControlFlowGraph* CFG){
    CFGNode* Node;
    
    LST_Rewind(CFG->NodeList);
    while((Node = LST_Next(CFG->NodeList))){
        Node->_Visited = 0;
        Node->TARJ_Index = 0;
        Node->TARJ_LowIndex = 0;
        Node->TARJ_HighestTransitionCount = 0;
        
        Node->SCC_Index = 0;
        Node->SCC_HighestTransitionCount = 0;
        LST_Destroy(Node->SCC_HotNodes);
        Node->SCC_HotNodes = NULL;
    }
}

static void CFG_ResetEdges(ControlFlowGraph *CFG, CFGNode* Node){
    CFGEdge *Edge;
    
    LST_Rewind(Node->Edges);
    
    Node->_Visited = 1;
    
    while((Edge = LST_Next(Node->Edges))){
        Edge->Count = 0;
        
        if(!Edge->Target->_Visited){
            CFG_ResetEdges(CFG, Edge->Target);
        }
    }
}

//Resets CFG execution data. Should be used only between epochs
void CFG_Reset(ControlFlowGraph *CFG){
    CFGNode *Node;
    
    LST_Rewind(CFG->NodeList);
    
    while((Node = LST_Next(CFG->NodeList))){
        Node->_Visited = 0;
        Node->EpochExecutionCount = 0;
    }
    
    LST_Rewind(CFG->NodeList);
    while((Node = LST_Next(CFG->NodeList))){
        if(!Node->_Visited)
            CFG_ResetEdges(CFG, Node);
    }    

    LST_Rewind(CFG->NodeList);
    while((Node = LST_Next(CFG->NodeList)))
        Node->_Visited = 0;
}

/*void CFG_ResetVisited(ControlFlowGraph *CFG){
    CFGNode *Node;
    
    LST_Rewind(CFG->NodeList);
    
    while((Node = LST_Next(CFG->NodeList))){
        Node->_Visited = 0;
    }
    
    CFG_ResetEdges(CFG, CFG_SearchNode(CFG, CFG->InitialAddress));    
}*/

static void PrintCFGEdges(ControlFlowGraph* CFG, CFGNode* Node, char* EdgesText){
    CFGEdge* Edge;
    char Scrap[128];
    unsigned char Alpha;
    
    LST_Rewind(Node->Edges);
    
    Node->_Visited = 1;
    
    while((Edge = LST_Next(Node->Edges))){
        
        //Only print edges which were used in the current epoch
        if(Edge->Count){
            
            Alpha = (Edge->Count > 200) ? 255 : (unsigned char)(((float)Edge->Count / 200) * 200) + 55;
            
            //Append edge
            sprintf(Scrap, "\t_%04X -> _%04X [color=\"#%02X0000\", label=\"%ld\"]\n", TU_GetStartingAddress(Node->Unit), TU_GetStartingAddress(Edge->Target->Unit), Alpha, Edge->Count);
            strcat(EdgesText, Scrap);
        }
        
        //If target has not been visited, visit it
        if(!(Edge->Target->_Visited)){
            PrintCFGEdges(CFG, Edge->Target, EdgesText);
        }
    }
}

void CFG_PrintCFG(ControlFlowGraph* CFG, char* DotFilename){
    
    char* NodeLabels;
    char* Scrap;
    char* EdgesText;
    FILE* DotFile;
    
    char* TUTYPE_BB =  "ellipse";
    char* TUTYPE_CFG = "box";
    
    char* NodeShape;
    
    CFGNode* Node;
    
    Scrap = malloc(256);
    NodeLabels = calloc(CFG->NumNodes, 64);
    EdgesText = calloc(1048576, 1);
    
    LST_Rewind(CFG->NodeList);
    while((Node = LST_Next(CFG->NodeList))){
        
        if(Node->EpochExecutionCount){
            
            NodeShape = (Node->Unit->Type == TUT_BB) ? TUTYPE_BB : TUTYPE_CFG;
            
            sprintf(Scrap, "\t_%04X [shape=%s, label=\"%04X\\n(%ld)\\n[%ld]\\n{%ld}\"]\n", TU_GetStartingAddress(Node->Unit), NodeShape, TU_GetStartingAddress(Node->Unit), Node->EpochExecutionCount, Node->SCC_Index, Node->Size);
            strcat(NodeLabels, Scrap);
        }
        
        Node->_Visited = 0;             // This should be a call to CFG_ResetCFGForTraversal.
    }
    
    LST_Rewind(CFG->NodeList);
    while((Node = LST_Next(CFG->NodeList))){
        if(!Node->_Visited)
            PrintCFGEdges(CFG, Node, EdgesText);
    }
    
    //PrintCFGEdges(CFG, CFG_SearchNode(CFG, CFG->InitialAddress), EdgesText);
    
    //printf(NodeLabels);
    //printf(EdgesText);
    
    sprintf(Scrap, "/home/juliano/Dropbox/grafos/%s", DotFilename);
    DotFile = fopen(Scrap, "w");
    
    sprintf(Scrap, "digraph CFG {\n\n");
    
    fprintf(DotFile, "%s", Scrap);
    fprintf(DotFile, "%s\n\n", NodeLabels);
    fprintf(DotFile, "%s\n\n", EdgesText);
    sprintf(Scrap, "}\n");
    
    fprintf(DotFile, "%s", Scrap);
    
    free(Scrap);
    free(NodeLabels);
    free(EdgesText);
    
    fclose(DotFile);
}

//Control Flow Graph Repository

ControlFlowGraph** CFGRepository;

void CFGRInit(void){
    CFGRepository = calloc(0x10000, sizeof(ControlFlowGraph*));
}

void CFGRRefresh(ControlFlowGraph* CFG){
    CFGNode* p;
    
    LST_Rewind(CFG->NodeList);
    
    while((p = LST_Next(CFG->NodeList)))
        CFGRNodeRefresh(CFG, p);
}

void CFGRBBRefresh(ControlFlowGraph* CFG, BasicBlock* BB){
    BBInstruction * p;
    
    LST_Rewind(BB->Instructions);
    
    while((p = LST_Next(BB->Instructions)))
        CFGRepository[p->Address] = CFG;
}

void CFGRAddressSet(ControlFlowGraph* CFG, uint16_t Address){
    CFGRepository[Address] = CFG;
}

void CFGRNodeRefresh(ControlFlowGraph* CFG, CFGNode* Node){
    TU_CFGRefresh(Node->Unit, CFG);
}

ControlFlowGraph* CFGRSearch(uint16_t Address){
    return CFGRepository[Address];
}
